import cv2
import numpy as np
import matplotlib.pyplot as plt

fileapple = 'teachers_code/apple.bmp'
filelion = 'teachers_code/lion.jpg'


# 灰度线性变换(将原来分布狭窄的灰度向两端拉开)
# scale = (255+2*delta)/255  t0=-delta
def hist_liner_trans(src, scale, t0):
    dst = src.astype(np.float32)
    dst = dst * scale + t0
    dst[dst > 255] = 255
    dst[dst < 0] = 0
    dst = dst.astype(np.uint8)
    return dst


def linerCorrection(filename):
    src = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
    # 灰度变换
    delta = 300
    scale = (255 + 2 * delta) / 255
    t0 = -delta + 50
    dst = hist_liner_trans(src, scale, t0)
    return src, dst


def histEqualize(filename):
    '图像均衡化（都是全局的均衡化）'
    src = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
    hist = cv2.calcHist([src], [0], None, [256], [0, 256]).reshape(-1)
    hist /= src.shape[0] * src.shape[1]

    lut = np.zeros(256, dtype=np.uint8)
    lut[0] = round(255 * hist[0])
    for i in range(1, 256):
        hist[i] += hist[i - 1]
        h = round(255 * hist[i - 1])
        if h > 255: h = 255
        lut[i] = h
    return src, cv2.LUT(src, lut)


def showResult(src, dst):
    ax1 = plt.subplot(221)
    ax1.set_title('Src Image')
    plt.imshow(src, cmap='gray', vmin=0, vmax=255)
    plt.xticks([])
    plt.yticks([])

    ax2 = plt.subplot(222)
    ax2.set_title('Src Histogram')
    plt.hist(src.ravel(), 256, [0, 256])

    # 显示目标图像和灰度直方图
    ax3 = plt.subplot(223)
    ax3.set_title('DST image')
    plt.imshow(dst, cmap='gray', vmin=0, vmax=255)
    plt.xticks([])
    plt.yticks([])

    ax2 = plt.subplot(224)
    ax2.set_title('DST Histogram')
    plt.hist(dst.ravel(), 256, [0, 256])

    plt.show()


def getSpecifyHist():
    dst_hist = np.zeros(256, dtype=np.float32)
    k = 2.0 / (255 ** 2)
    for i in range(256):
        dst_hist[i] = k * i
    return dst_hist


def histSpecify(filename, dst_hist):
    '直方图规定化'
    src = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
    # 读取原图的直方图
    src_hist = cv2.calcHist([src], [0], None, [256], [0, 256]).flatten()
    # 对直方图信息进行归一化处理，结果为每个灰度值的概率密度
    # shape（0）代表图像垂直尺寸，shape（1）代表图像水平尺寸
    src_hist /= src.shape[0] * src.shape[1]

    src_cdf = np.zeros(256, dtype=np.float32)
    dst_cdf = np.zeros(256, dtype=np.float32)
    src_cdf[0] = src_hist[0]
    dst_cdf[0] = dst_hist[0]
    for i in range(1, 256):
        # 计算原图像的累计分布函数的值（为取当前值的概率密度加取前面所有值的概率密度）
        src_cdf[i] = src_hist[i] + src_cdf[i - 1]
        # 计算规定化直方图的累计分布函数的值（为取当前值的概率密度加取前面所有值的概率密度）
        dst_cdf[i] = dst_hist[i] + dst_cdf[i - 1]

    lut = np.zeros(256, dtype=np.uint8)
    for i in range(0, 256):
        # 寻求累积分布函数最接近的原图像灰度i, 和参考灰度j
        # 这等价于用目标分布函数与i所对应的累积分布函数值求差, 再寻找差的最小值做对应的j
        diff_cdf = np.abs(dst_cdf - src_cdf[i])
        # 选取diff_cdf中的最小值
        j = np.argmin(diff_cdf)
        lut[i] = j
    # LUT函数：
    # 例：lut[1] 为20，
    # 意思为：src中灰度为1的所有像素的灰度变为20
    return src, cv2.LUT(src, lut)


def test():
    # src, dst = histEqualize(filelion)
    src, dst = histSpecify(filelion, getSpecifyHist())
    # src, dst = linerCorrection(filelion)
    showResult(src, dst)


test()
